import numpy as np  # 导入numpy工具包
from os import listdir  # 使用listdir模块，用于访问本地文件
from sklearn import neighbors
import load_mnist
import time

if __name__ == '__main__':
    path = r'E:\python\python代码\dataSet'

    train_num = 6000  # 训练集数量     上限为6W
    test_num = 600  # 测试集数量     上限为1W

    # read dataSet
    train_dataSet, train_hwLabels = load_mnist.load_mnist_train(path)
    knn = neighbors.KNeighborsClassifier(algorithm='kd_tree', n_neighbors=3)
    print("数据读取完成，开始训练")
    # print(train_dataSet[:5])
    # print(train_hwLabels[:5])
    T1 = time.time()
    knn.fit(train_dataSet[:train_num], train_hwLabels[:train_num])

    # read  testing dataSet
    dataSet, hwLabels = load_mnist.load_mnist_test(path)

    res = knn.predict(dataSet[:test_num])  # 对测试集进行预测
    error_num = np.sum(res != hwLabels[:test_num])  # 统计分类错误的数目

    T2 = time.time()
    print("测试集数量：", test_num, " 错误次数：", error_num, "  错误率：", error_num / float(test_num))
    print("总共用时：", T2 - T1)
